package com.example.algorithm.tree.uf;

public class UF_TREE_Weighted {

    /**
     * 记录结点元素和该元素所在分组的标识
     */
    private int[] eleAndGroup;

    /**
     * 记录并查集中数据的分组个数
     */
    private int count;

    /**
     * 用来存储每个根节点对应的树中结点的个数
     */
    private int[] sz;

    /**
     * 初始化并查集数
     * @param N
     */
    public UF_TREE_Weighted(int N){
        //初始化分组的数量,默认情况下有N个
        this.count = N;

        //初始化数组
        this.eleAndGroup = new int[N];

        //初始化数组的元素及所在组的标识符号,让eleAndGroup的索引作为并查集每个结点的元素，并且让每个索引处的值(该元素所在的组的标识符)就是该索引的分组
        for (int i = 0; i < eleAndGroup.length; i++) {
            eleAndGroup[i] = i;
        }
        //默认每个sz都是1
        this.sz = new int[N];
        for (int i = 0; i < sz.length; i++) {
            sz[i] = 1;
        }
    }

    /**
     * 获取当前并查集的数据有多少个分组
     * @return
     */
    public int count(){
        return count;
    }

    /**
     * 判断p和q是否在一个分组中
     */
    public boolean connected(int p,int q){
        return find(p) == find(q);
    }


    /**
     * 元素p所在分组的标识符
     * @param p
     * @return
     */
    public int find(int p){

        while (true){
            if (p == eleAndGroup[p]){
                return p;
            }
            p = eleAndGroup[p];
        }
    }

    /**
     *  把p元素所在的分组和q所在的分组合并
     */
    public void union(int p, int q){
        //找到p元素和q元素所在组的根节点
        int pRoot = find(p);
        int qRoot = find(q);

        //如果p和q所在根节点一样，则不需要合并
        if (pRoot == qRoot){
            return;
        }

        //比较proot对应树的大小和qroot对应树的大小，将小树放到大树中
        if (sz[pRoot] < sz[qRoot]){
            eleAndGroup[pRoot] = qRoot;
            sz[qRoot] += sz[qRoot];
        }else {
            eleAndGroup[qRoot] = pRoot;
            sz[pRoot] += sz[qRoot];
        }




        this.count --;
    }


}
